package src;

import org.apache.poi.openxml4j.exceptions.InvalidFormatException;
import org.apache.poi.util.Units;
import org.apache.poi.xwpf.usermodel.*;
import org.openxmlformats.schemas.drawingml.x2006.picture.CTPicture;
import org.openxmlformats.schemas.officeDocument.x2006.math.CTOMath;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.*;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.*;

public class WordDocumentExample {
    public static void main(String[] args) throws Exception {
        String sourceFilePath = "src/main/resources/word/index.docx"; // 源文档路径
        String outputFilePath = "src/main/resources/word/unpack/"; // 拆分后的文件输出目录


        FileInputStream fis = new FileInputStream(sourceFilePath);
        XWPFDocument sourceDocument = new XWPFDocument(fis);
        Map<String, List<IBodyElement>> separation = separation(sourceDocument);
        unpack(outputFilePath,separation);

    }

    public static Map<String,List<IBodyElement>> separation(XWPFDocument sourceDocument){
        Map<String,List<IBodyElement>> map = new HashMap<>();
        List<IBodyElement> bodyElements = sourceDocument.getBodyElements();
        String curr = "";
        for (int i = 0; i < bodyElements.size(); i++) {
            IBodyElement bodyElement = sourceDocument.getBodyElements().get(i);

            if (bodyElement instanceof XWPFTable){
                if(!map.containsKey(curr))map.put(curr,new ArrayList<>()); // 添加
                map.get(curr).add(bodyElement);
            }else if (bodyElement instanceof XWPFParagraph){
                XWPFParagraph paragraph = (XWPFParagraph) bodyElement;

                if (isHeading(paragraph)) {
                    //curr = paragraphs.get(i).getText().replaceAll("^\\d+(\\.\\d+)*\\s*", "");
                    curr = paragraph.getText();
                } else {
                    if(!map.containsKey(curr))map.put(curr,new ArrayList<>()); // 添加
                    map.get(curr).add(bodyElement);
                }
            }
        }

        return map;
    }

    public static void unpack(String path,Map<String,List<IBodyElement>> map){
        new File(path).mkdirs();
        map.forEach((key, vars) ->{
            XWPFDocument document = new XWPFDocument();

            vars.forEach(var -> {

                if (var instanceof XWPFParagraph) {
                    XWPFParagraph newParagraph = document.createParagraph();
                    XWPFParagraph paragraph = (XWPFParagraph)var;

                    for (IRunElement iRun : paragraph.getIRuns()) {

                        if(iRun instanceof XWPFRun){
                            XWPFRun run = (XWPFRun)iRun;

                            List<XWPFPicture> pictures = run.getEmbeddedPictures();

                            if (pictures.size() == 0) {
                                newParagraph.getCTP().set(paragraph.getCTP());
                                newParagraph.getCTPPr().set(paragraph.getCTPPr());
                            }
                            for (XWPFPicture picture : pictures) {
                                XWPFPictureData pictureData = picture.getPictureData();
                                byte[] pictureBytes = pictureData.getData();
                                int pictureType = pictureData.getPictureType();
                                String fileName = pictureData.getFileName();

                                try {
                                    XWPFRun newRun = newParagraph.createRun();
                                    newRun.addPicture(new ByteArrayInputStream(pictureBytes), pictureType, fileName, Units.toEMU(picture.getWidth()), Units.toEMU(picture.getDepth()));
                                } catch (Exception e) {
                                    throw new RuntimeException(e);
                                }

                            }
                        }
                    }

                } else if (var instanceof XWPFTable) {
                    XWPFTable table = (XWPFTable)var;
                    XWPFTable newTable = document.createTable();
                    newTable.getCTTbl().set(table.getCTTbl());
                }


            });


            try (FileOutputStream out = new FileOutputStream(path+key+".docx")) {
                document.write(out);
            } catch (IOException e) {
                e.printStackTrace();
            }

        });
    }

    private static boolean isHeading(XWPFParagraph paragraph) {
        if (paragraph == null) return false;
        return (paragraph.getStyle() != null &&
                (paragraph.getStyle().equals("1") ||
                        paragraph.getStyle().equals("2") ||
                        paragraph.getStyle().equals("3") ||
                        paragraph.getStyle().equals("4")));
    }
}


